/*
 * Copyright 2020 Netflix, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.netflix.graphql.dgs.subscriptions.websockets

import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import com.netflix.graphql.dgs.DgsQueryExecutor
import graphql.ExecutionResult
import io.mockk.*
import io.mockk.impl.annotations.MockK
import io.mockk.junit5.MockKExtension
import io.reactivex.rxjava3.core.Flowable
import org.assertj.core.api.Assertions.assertThat
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.extension.ExtendWith
import org.reactivestreams.Publisher
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketSession

@ExtendWith(MockKExtension::class)
class DgsWebsocketHandlerTest {

    private lateinit var dgsWebsocketHandler: DgsWebSocketHandler

    @BeforeEach
    fun setup() {
        dgsWebsocketHandler = DgsWebSocketHandler(dgsQueryExecutor)

        every { session1.id } returns "1"
        every { session2.id } returns "2"
    }

    @MockK
    lateinit var dgsQueryExecutor: DgsQueryExecutor

    @MockK
    lateinit var session1: WebSocketSession

    @MockK
    lateinit var session2: WebSocketSession

    @MockK
    lateinit var executionResult: ExecutionResult

    private val queryMessage = TextMessage("""{
                "type": "$GQL_START",
                "payload": {
                   "query": "{ hello }",
                   "variables": {},
                   "extensions": {}
                }
            }
            """.trimIndent())

    @Test
    fun testMultipleClients() {
        connect(session1)
        connect(session2)
        start(session2, 3)
        start(session1, 1)

        disconnect(session2)
        disconnect(session1)

        // ACK, DATA, COMPLETE
        verify(exactly = 3) {
            session1.sendMessage(any())
        }

        // ACK, DATA, DATA DATA, COMPLETE
        verify(exactly = 5) {
            session2.sendMessage(any())
        }
    }

    @Test
    fun testWithError() {
        connect(session1)
        startWithError(session1)
        disconnect(session1)

        // ACK, ERROR
        verify(exactly = 2) {
            session1.sendMessage(any())
        }
    }

    private fun connect(webSocketSession: WebSocketSession) {
        val currentNrOfSessions = dgsWebsocketHandler.sessions.size

        val slot = slot<TextMessage>()
        every { webSocketSession.sendMessage(capture(slot)) } just Runs

        val textMessage = TextMessage("""{
                "type": "$GQL_CONNECTION_INIT"
            }
            """.trimIndent())

        dgsWebsocketHandler.handleTextMessage(webSocketSession, textMessage)

        assertThat(dgsWebsocketHandler.sessions.size).isEqualTo(currentNrOfSessions + 1)

        val returnMessage = jacksonObjectMapper().readValue<OperationMessage>(slot.captured.asBytes())
        assertThat(returnMessage.type).isEqualTo(GQL_CONNECTION_ACK)
    }

    private fun disconnect(webSocketSession: WebSocketSession) {
        val currentNrOfSessions = dgsWebsocketHandler.sessions.size
        every { webSocketSession.close() } just Runs

        val textMessage = TextMessage("""{
                "type": "$GQL_CONNECTION_TERMINATE"
            }
            """.trimIndent())

        dgsWebsocketHandler.handleTextMessage(webSocketSession, textMessage)

        assertThat(dgsWebsocketHandler.sessions.size).isEqualTo(currentNrOfSessions - 1)

        verify { webSocketSession.close() }
    }

    private fun start(webSocketSession: WebSocketSession, nrOfResults: Int) {


        every { webSocketSession.isOpen } returns true

        val results = (1..nrOfResults).map {
            val result1 = mockkClass(ExecutionResult::class)
            every { result1.getData<Any>()} returns it
            result1
        }

        every { executionResult.getData<Publisher<ExecutionResult>>() } returns Flowable.fromIterable(results)
        every { dgsQueryExecutor.execute("{ hello }") } returns executionResult

        dgsWebsocketHandler.handleTextMessage(webSocketSession, queryMessage)


    }

    private fun startWithError(webSocketSession: WebSocketSession) {
        every { webSocketSession.isOpen } returns true
        every { executionResult.getData<Publisher<ExecutionResult>>() } returns Flowable.error(RuntimeException("That's wrong!"))
        every { dgsQueryExecutor.execute("{ hello }") } returns executionResult

        dgsWebsocketHandler.handleTextMessage(webSocketSession, queryMessage)
    }
}